import json
import os
import warnings
import re
# import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator

import numpy as np
# data_file_path = "./data/localhost.localdomain-Comm_NUMAMemcpy_GPUToHost-numa0-cuda1.json"
folder = "F:/Codes/sugon_code/xgmi-scope-plot/data"
# dirs= os.listdir(folder)
# benchmarks_lists = []
# json_dict = {}




def generate_dict(folder):
    json_dict = {}
    json_files = []
    for root,dirs,files in os.walk(folder):
        for name in files:
            # print(os.path.join(root, name))
            json_files .append(os.path.join(root, name))
        # for name in dirs:
        #     print(os.path.join(root, name))
    category_detail_set = set()
    iteration_benchmark = []
    mean_benchmark = []
    stddev_benchmark = []
    median_benchmark = []
    pattern1 = re.compile(r':([\d]+)')
    pattern2 = re.compile(r'([\d]+)')
    for json_file in json_files:
        # print(json_file)
        size = os.path.getsize(json_file)
        if size == 0:
            # warnings.warn("There is nothing in :"+json_file)
            continue
        with open(json_file,"r") as load_f:
            load_dict = json.load(load_f)
            benchmark_length = len(load_dict['benchmarks'])
            # print(json_file)
            # print(len(load_dict['benchmarks']))
            for benchmark in load_dict['benchmarks']:
                str1 = pattern1.search(benchmark['name'])[0]
                logn = pattern2.search(str1)[0]
                category = benchmark['name'].split('/')[0]
                category_detail =  benchmark['name'].split('log2')[0]
                benchmark['name'].split('log2')[0].split('/')[-2]
                benchmark['logn'] = int(logn.strip())
                benchmark['category'] = category.strip()
                benchmark['category_detail'] = category_detail.strip()
                benchmark['GB/s']=int(benchmark['bytes_per_second'])/(1024*1024*1024)
                # benchmark['B'] = r'$2^%d$' % benchmark['logn']
                devs = benchmark['name'].split('log2')[0].split('/')[1:-1]
                # print(devs)
                benchmark['cuda_id'] = int(devs[1]) 
                benchmark['numa_id'] = int(devs[0])
                if(len(devs)==3):
                    benchmark['cuda2_id'] = int(devs[2])
                # print(devs)
                # print(benchmark['cuda_id'],benchmark['numa_id'])
                
                # print(category)
                # print(category_detail)
                if benchmark['run_type'] == 'iteration':
                    iteration_benchmark.append(benchmark)
                if benchmark['run_type'] == 'aggregate':
                    if benchmark['aggregate_name'] == 'mean':
                        mean_benchmark.append(benchmark)
                    if benchmark['aggregate_name'] == 'median':
                        median_benchmark.append(benchmark)
                    if benchmark['aggregate_name'] == 'stddev':
                        stddev_benchmark.append(benchmark)
            json_dict['mean_benchmark'] = mean_benchmark
            json_dict['median_benchmark'] = median_benchmark
            json_dict['stddev_benchmark'] = stddev_benchmark
    return json_dict

def plot(json_dict,copy_name,type_='stddev',host_name="",cuda_id=None, numa_id=None,x_axis="logn",y_axis = 'GB/s',low=8,high=33,need_xlabel=True,need_ylabel=True,x_alias='Transfer Size (B)',y_alias=None):
    data_base = []
    if type_ == 'stddev':
        data_base = json_dict['stddev_benchmark']
    elif type_ == 'mean':
        data_base = json_dict['mean_benchmark']
    elif type_ == 'median':
        data_base = json_dict['median_benchmark']
    else :
        type_ = None
        assert(type_!=None)
        # type类型不存在
    final_base = data_base.copy()
    category_detail_set=set()
    category_detail_list = []
    ## 筛选
    for benchmark in data_base:
        # print(copy_name,"--",benchmark['category'])
        # print(len(final_base))
        # print(final_base)
        if copy_name not in benchmark['category']:
            # print("delete")
            final_base.remove(benchmark)
            
        elif cuda_id != None and cuda_id !=benchmark['cuda_id']:
            final_base.remove(benchmark)
        elif numa_id != None and numa_id !=benchmark['numa_id']:
            final_base.remove(benchmark)
        elif benchmark['logn'] < low or benchmark['logn'] > high:
            final_base.remove(benchmark)
        else:
            category_detail_ = benchmark["category_detail"]
            category_detail_list.append(category_detail_)
            category_detail_set.add(category_detail_)
    category_detail_set = set(category_detail_list)
    # print(category_detail_set)
    # sns.relplot(x=x_axis,y=y_axis,kind="line",data=final_base)
    # fig, ax = plt.subplots()  # Create a figure and an axes.
    # print(len(category_detail_set))
    ## 对set中的字符串排序
    category_detail_list = list(category_detail_set)
    category_detail_list = sorted(category_detail_list)
    biggest = 0
    for category_detail_ in category_detail_list:
        # print(category_detail_)
        x_data = []
        y_data = []
        for benchmark in final_base:
            if category_detail_ in benchmark["name"]:
                x_data.append(benchmark[x_axis])
                y_data.append(benchmark[y_axis]) 
                if benchmark[y_axis] > biggest:
                    biggest = benchmark[y_axis] 
        devs = category_detail_.split('log2')[0].split('/')[1:-1]
        # print(devs)
        if(len(devs)==2):
            label_ = "numa_id:"+str(devs[0]) + ",cuda_id:"+str(devs[1])    
        elif(len(devs)==3):
            label_ =  "numa_id:"+str(devs[0]) + ",cuda_id:"+str(devs[1])+ ",cuda2_id:"+str(devs[2])

        # print(x_data)
        # print(y_data)
        zipped = zip(x_data,y_data)
        sort_zipped = sorted(zipped,key=lambda x:(x[0],x[1]))
        result = zip(*sort_zipped)
        x_data, y_data = [list(x) for x in result]
        # print(x_axis)
        # print(y_axis)

        plt.plot(x_data,y_data,label=label_,marker='|')
        # plt.show()
    if need_xlabel:
        if x_alias == None:
            plt.xlabel(x_axis)
        else:
            plt.xlabel(x_alias)
    if need_ylabel:
        if y_alias == None:
            plt.ylabel(y_axis)
        else:
            plt.ylabel(y_alias)
    #x轴是唯一的，生成X轴
    x_data_sci=[]
    for i in x_data:
        x_data_sci.append(r'${2}^{%d}$'%i)
    x_major_locator=MultipleLocator(4)
    #把x轴的刻度间隔设置为1，并存在变量里
    y_major_locator=MultipleLocator(5)
    #把y轴的刻度间隔设置为10，并存在变量里
    plt.ylim([0,biggest+5])
    plt.xlim([low,high])
    ax=plt.gca()
    #ax为两条坐标轴的实例
    ax.xaxis.set_major_locator(x_major_locator)
    #把x轴的主刻度设置为1的倍数
    ax.yaxis.set_major_locator(y_major_locator)
    #把y轴的主刻度设置为10的倍数
    # ax.set_xticklabels(x_data_sci)
    
    plt.xticks(x_data[2:24:4],x_data_sci[2:24:4])
    plt.title(host_name+copy_name)  # Add a title to the axes.
    plt.legend()  # Add a legend
    plt.grid()
    


ax1 = plt.subplot(211)
json_dict = generate_dict(folder)
plot(json_dict,"Comm_NUMAMemcpy_PinnedToGPU",low = 8,high = 33,type_='mean',need_xlabel=False,y_alias='Transfer Bandwidth(GB/s)')
ax1.set_xticklabels([])



# ax2 = plt.subplot(212,sharex=ax1)
ax2 = plt.subplot(212)
json_dict = generate_dict(folder)
plot(json_dict,"Comm_NUMAMemcpy_GPUToHost",low = 8,high = 33,type_='mean',y_alias='Transfer Bandwidth(GB/s)')
# plt.subplots_adjust(wspace=1,hspace=1)
ax1.get_shared_y_axes().join(ax1, ax2)

plt.show()

    

# import json  
# f =open(data_file_path,encoding='utf-8') #打开‘product.json’的json文件
# res=f.read()  #读文件
# print(json.loads(res))#把json串变成python的数据类型：字典      
